#/usr/bin/python
# -*- coding: utf-8 -*-
'''

...

Created on 6 mars 2012

@author: Ludovic Barrière and Steven Marguet, Institut Clément Ader, 135 av. de Rangueil, 31077 Toulouse, France

© copyright:
    
    TODO
'''

import mpi4py.MPI as mpi

import scipy.sparse as sp
import scipy.sparse.linalg as spla
import numpy as np

## Parent class of linear system solvers:
class LinearSystemSolver(object): # abstract class
    
    def __init__(self):
        self.A = None
        self.b = None
        self.x = None
    
    def solve(self,A,b,output=sp.lil_matrix):
        raise NotImplementedError

## Child classes of linear system solvers:   
class Direct(LinearSystemSolver):
    
    def __init__(self):
        
        super(Direct,self).__init__()
        
        
    def solve(self,A,b,output=sp.lil_matrix):
        self.x = output(spla.spsolve(A,b))
        return self.x

   
class IterativeSolver(LinearSystemSolver): # abstract class
    
    def __init__(self):
        super(IterativeSolver,self).__init__()
        self.preconditioner = None
        self.constraint = None

    def setPreconditioner(self,M):
        self.preconditioner = M
    
    def setConstraint(self,H):
        self.constraint = H

## Child classes of iterative solvers:
class KrylovSolver(IterativeSolver): # abstract class
    
    def __init__(self):
        super.__init__(self)
        self.KrylovSubspace = None
         
    
## Child classes of KrylovSolvers:
class GMRES(KrylovSolver):
    
    def __init__(self):
        super.__init__(self)
        
    def solve(self,A,b,M=None,H=None):
        pass
        
class CG(KrylovSolver):

    def __init__(self):
        super.__init__(self)
        
    def solve(self,A,b,M=None,H=None):
        pass


class ParallelGMRES(KrylovSolver):
    '''
        
    '''
    def __init__(self):
        super.__init__(self)
        
    def solve(self, master, lhslist, rhslist):
        '''
            Solve the linear system.
        '''
        
        pass
    
    @staticmethod
    def localStep(worker, lhslist, rhslist):
        '''
           Local step of the parallelized linear system solver (GMRES). 
        '''
        pass
    
        
class ParallelCG(KrylovSolver):
    '''
        Solving the linear system (Ax=b) assuming that the master object has workers. 
        Workers are supposed to hold their own operators such that $A = \sum_{i} A_{i}$ and
        $b = \sum_{i} bi$.
    '''
    def __init__(self):
        super.__init__(self)
        
    def solve(self, master, lhslist, rhslist):
        '''
            Solve the linear system.
        '''
        
        # initialize the local process: 
        master.workers.bcast('staticmethod',[self.localStep])
    
        # manage the conjugate gradient:
        
    
    @staticmethod
    def localStep(worker, lhslist, rhslist):
        '''
           Local step of the parallelized linear system solver (CG).
        '''
        Ai = worker.getLhs(lhslist, full=False)
        bi = worker.getRhs(rhslist, full=False)
        next = 1
        
        while next:
            # stop/continue message
            next = worker.master.bcast(root=0)
            if not next:
                break
            
            
            
        
        
        

